// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/function.h>

#include <cstddef>
#include <memory>
#include <sstream>

#include <nebula/compute/api_scalar.h>
#include <nebula/compute/cast.h>
#include <nebula/compute/exec.h>
#include <nebula/compute/exec_internal.h>
#include <nebula/compute/function_internal.h>
#include <nebula/compute/function_options.h>
#include <nebula/compute/kernels/common_internal.h>
#include <nebula/compute/registry.h>
#include <nebula/core/datum.h>
#include <nebula/util/cpu_info.h>
#include <turbo/log/logging.h>
#include <nebula/trace/tracing_internal.h>

namespace nebula::compute {

    turbo::Result<std::shared_ptr<Buffer>> FunctionOptionsType::serialize(
            const FunctionOptions &) const {
        return turbo::unimplemented_error("serialize for ", type_name());
    }

    turbo::Result<std::unique_ptr<FunctionOptions>> FunctionOptionsType::deserialize(
            BufferSpan buffer) const {
        return turbo::unimplemented_error("deserialize for ", type_name());
    }

    std::string FunctionOptions::to_string() const { return options_type()->stringify(*this); }

    bool FunctionOptions::equals(const FunctionOptions &other) const {
        if (this == &other) return true;
        if (options_type() != other.options_type()) return false;
        return options_type()->compare(*this, other);
    }

    std::unique_ptr<FunctionOptions> FunctionOptions::copy() const {
        return options_type()->copy(*this);
    }

    turbo::Result<std::shared_ptr<Buffer>> FunctionOptions::serialize() const {
        return options_type()->serialize(*this);
    }

    turbo::Result<std::unique_ptr<FunctionOptions>> FunctionOptions::deserialize(
            const std::string &type_name, BufferSpan buffer) {
        TURBO_MOVE_OR_RAISE(auto options,
                            get_function_registry()->get_function_options_type(type_name));
        return options->deserialize(buffer);
    }

    void print_to(const FunctionOptions &options, std::ostream *os) {
        *os << options.to_string();
    }

    static const FunctionDoc kEmptyFunctionDoc{};

    const FunctionDoc &FunctionDoc::Empty() { return kEmptyFunctionDoc; }

    static turbo::Status CheckArityImpl(const Function &func, int num_args) {
        if (func.arity().is_varargs && num_args < func.arity().num_args) {
            return turbo::invalid_argument_error("VarArgs function '", func.name(), "' needs at least ",
                                                 func.arity().num_args, " arguments but only ", num_args,
                                                 " passed");
        }

        if (!func.arity().is_varargs && num_args != func.arity().num_args) {
            return turbo::invalid_argument_error("Function '", func.name(), "' accepts ", func.arity().num_args,
                                                 " arguments but ", num_args, " passed");
        }
        return turbo::OkStatus();
    }

    turbo::Status Function::check_arity(size_t num_args) const {
        return CheckArityImpl(*this, static_cast<int>(num_args));
    }

    namespace {

        turbo::Status CheckOptions(const Function &function, const FunctionOptions *options) {
            if (options == nullptr && function.doc().options_required) {
                return turbo::invalid_argument_error("Function '", function.name(),
                                                     "' cannot be called without options");
            }
            return turbo::OkStatus();
        }

    }  // namespace

    namespace detail {

        turbo::Status no_matching_kernel(const Function *func, const std::vector<TypeHolder> &types) {
            return turbo::unimplemented_error("Function '", func->name(),
                                              "' has no kernel matching input types ",
                                              TypeHolder::to_string(types));
        }

        template<typename KernelType>
        const KernelType *dispatch_exact_impl(const std::vector<KernelType *> &kernels,
                                            const std::vector<TypeHolder> &values) {
            const KernelType *kernel_matches[SimdLevel::MAX] = {nullptr};

            // Validate arity
            for (const auto &kernel: kernels) {
                if (kernel->signature->matches_inputs(values)) {
                    kernel_matches[kernel->simd_level] = kernel;
                }
            }

            // Dispatch as the CPU feature
#if defined(NEBULA_HAVE_RUNTIME_AVX512) || NEBULA_HAVE_RUNTIME_AVX2_SUPPORTED
            auto cpu_info = nebula::internal::CpuInfo::GetInstance();
#endif
#if defined(NEBULA_HAVE_RUNTIME_AVX512)
            if (cpu_info->is_supported(nebula::internal::CpuInfo::AVX512)) {
              if (kernel_matches[SimdLevel::AVX512]) {
                return kernel_matches[SimdLevel::AVX512];
              }
            }
#endif
#if NEBULA_HAVE_RUNTIME_AVX2_SUPPORTED
            if (cpu_info->is_supported(nebula::internal::CpuInfo::AVX2)) {
                if (kernel_matches[SimdLevel::AVX2]) {
                    return kernel_matches[SimdLevel::AVX2];
                }
            }
#endif
            if (kernel_matches[SimdLevel::NONE]) {
                return kernel_matches[SimdLevel::NONE];
            }

            return nullptr;
        }

        const Kernel *dispatch_exact_impl(const Function *func,
                                        const std::vector<TypeHolder> &values) {
            if (func->kind() == Function::SCALAR) {
                return dispatch_exact_impl(turbo::checked_cast<const ScalarFunction *>(func)->kernels(),
                                         values);
            }

            if (func->kind() == Function::VECTOR) {
                return dispatch_exact_impl(turbo::checked_cast<const VectorFunction *>(func)->kernels(),
                                         values);
            }

            if (func->kind() == Function::SCALAR_AGGREGATE) {
                return dispatch_exact_impl(
                        turbo::checked_cast<const ScalarAggregateFunction *>(func)->kernels(), values);
            }

            if (func->kind() == Function::HASH_AGGREGATE) {
                return dispatch_exact_impl(turbo::checked_cast<const HashAggregateFunction *>(func)->kernels(),
                                         values);
            }

            return nullptr;
        }

        struct FunctionExecutorImpl : public FunctionExecutor {
            FunctionExecutorImpl(std::vector<TypeHolder> in_types, const Kernel *kernel,
                                 std::unique_ptr<detail::KernelExecutor> executor,
                                 const Function &func)
                    : in_types(std::move(in_types)),
                      kernel(kernel),
                      kernel_ctx(default_exec_context(), kernel),
                      executor(std::move(executor)),
                      func(func),
                      state(),
                      options(nullptr),
                      inited(false) {}

            virtual ~FunctionExecutorImpl() {}

            turbo::Status KernelInit(const FunctionOptions *options) {
                TURBO_RETURN_NOT_OK(CheckOptions(func, options));
                if (options == nullptr) {
                    options = func.default_options();
                }
                if (kernel->init) {
                    TURBO_MOVE_OR_RAISE(state,
                                        kernel->init(&kernel_ctx, {kernel, in_types, options}));
                    kernel_ctx.set_state(state.get());
                }

                TURBO_RETURN_NOT_OK(executor->init(&kernel_ctx, {kernel, in_types, options}));
                this->options = options;
                inited = true;
                return turbo::OkStatus();
            }

            turbo::Status init(const FunctionOptions *options, ExecContext *exec_ctx) override {
                if (exec_ctx == nullptr) {
                    exec_ctx = default_exec_context();
                }
                kernel_ctx = KernelContext{exec_ctx, kernel};
                return KernelInit(options);
            }

            turbo::Result<Datum> execute(const std::vector<Datum> &args, int64_t passed_length) override {
                util::tracing::Span span;

                auto func_kind = func.kind();
                const auto &func_name = func.name();
                START_COMPUTE_SPAN(span, func_name,
                                   {
                                       { "function.name", func_name },
                                       { "function.options", options ? options->to_string() : "<nullptr>" },
                                       { "function.kind", func_kind }
                                   });

                if (in_types.size() != args.size()) {
                    return turbo::invalid_argument_error("Execution of '", func_name, "' expected ", in_types.size(),
                                                         " arguments but got ", args.size());
                }
                if (!inited) {
                    TURBO_RETURN_NOT_OK(init(nullptr, default_exec_context()));
                }
                ExecContext *ctx = kernel_ctx.exec_context();
                // Cast arguments if necessary
                std::vector<Datum> args_with_cast(args.size());
                for (size_t i = 0; i != args.size(); ++i) {
                    const auto &in_type = in_types[i];
                    auto arg = args[i];
                    if (in_type != args[i].type()) {
                        TURBO_MOVE_OR_RAISE(arg, Cast(args[i], CastOptions::Safe(in_type), ctx));
                    }
                    args_with_cast[i] = std::move(arg);
                }

                detail::DatumAccumulator listener;

                ExecBatch input(std::move(args_with_cast), /*length=*/0);
                if (input.num_values() == 0) {
                    if (passed_length != -1) {
                        input.length = passed_length;
                    }
                } else {
                    bool all_same_length = false;
                    int64_t inferred_length = detail::InferBatchLength(input.values, &all_same_length);
                    input.length = inferred_length;
                    if (func_kind == Function::SCALAR) {
                        if (passed_length != -1 && passed_length != inferred_length) {
                            return turbo::invalid_argument_error(
                                    "Passed batch length for execution did not match actual"
                                    " length of values for execution of scalar function '",
                                    func_name, "'");
                        }
                    } else if (func_kind == Function::VECTOR) {
                        auto vkernel = static_cast<const VectorKernel *>(kernel);
                        if (!all_same_length && vkernel->can_execute_chunkwise) {
                            return turbo::invalid_argument_error("Arguments for execution of vector kernel function '",
                                                                 func_name, "' must all be the same length");
                        }
                    }
                }
                TURBO_RETURN_NOT_OK(executor->execute(input, &listener));
                const auto out = executor->wrap_results(input.values, listener.values());
#ifndef NDEBUG
                        KCHECK_OK(executor->CheckResultType(out, func_name.c_str()));
#endif
                return out;
            }

            std::vector<TypeHolder> in_types;
            const Kernel *kernel;
            KernelContext kernel_ctx;
            std::unique_ptr<detail::KernelExecutor> executor;
            const Function &func;
            std::unique_ptr<KernelState> state;
            const FunctionOptions *options;
            bool inited;
        };

    }  // namespace detail

    turbo::Result<const Kernel *> Function::dispatch_exact(
            const std::vector<TypeHolder> &values) const {
        if (kind_ == Function::META) {
            return turbo::unimplemented_error("Dispatch for a MetaFunction's Kernels");
        }
        TURBO_RETURN_NOT_OK(check_arity(values.size()));

        if (auto kernel = detail::dispatch_exact_impl(this, values)) {
            return kernel;
        }
        return detail::no_matching_kernel(this, values);
    }

    turbo::Result<const Kernel *> Function::dispatch_best(std::vector<TypeHolder> *values) const {
        // TODO(ARROW-11508) permit generic conversions here
        return dispatch_exact(*values);
    }

    turbo::Result<std::shared_ptr<FunctionExecutor>> Function::get_best_executor(
            std::vector<TypeHolder> inputs) const {
        std::unique_ptr<detail::KernelExecutor> executor;
        if (kind() == Function::SCALAR) {
            executor = detail::KernelExecutor::MakeScalar();
        } else if (kind() == Function::VECTOR) {
            executor = detail::KernelExecutor::MakeVector();
        } else if (kind() == Function::SCALAR_AGGREGATE) {
            executor = detail::KernelExecutor::MakeScalarAggregate();
        } else {
            return turbo::unimplemented_error("Direct execution of HASH_AGGREGATE functions");
        }

        TURBO_MOVE_OR_RAISE(const Kernel *kernel, dispatch_best(&inputs));

        return std::make_shared<detail::FunctionExecutorImpl>(std::move(inputs), kernel,
                                                              std::move(executor), *this);
    }

    namespace {

        turbo::Result<Datum> ExecuteInternal(const Function &func, std::vector<Datum> args,
                                             int64_t passed_length, const FunctionOptions *options,
                                             ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(auto inputs, internal::GetFunctionArgumentTypes(args));
            TURBO_MOVE_OR_RAISE(auto func_exec, func.get_best_executor(inputs));
            TURBO_RETURN_NOT_OK(func_exec->init(options, ctx));
            return func_exec->execute(args, passed_length);
        }

    }  // namespace

    turbo::Result<Datum> Function::execute(const std::vector<Datum> &args,
                                           const FunctionOptions *options, ExecContext *ctx) const {
        return ExecuteInternal(*this, args, /*passed_length=*/-1, options, ctx);
    }

    turbo::Result<Datum> Function::execute(const ExecBatch &batch, const FunctionOptions *options,
                                           ExecContext *ctx) const {
        return ExecuteInternal(*this, batch.values, batch.length, options, ctx);
    }

    namespace {

        turbo::Status ValidateFunctionSummary(const std::string &s) {
            if (s.find('\n') != s.npos) {
                return turbo::invalid_argument_error("summary contains a newline");
            }
            if (s.back() == '.') {
                return turbo::invalid_argument_error("summary ends with a point");
            }
            return turbo::OkStatus();
        }

        turbo::Status ValidateFunctionDescription(const std::string &s) {
            if (!s.empty() && s.back() == '\n') {
                return turbo::invalid_argument_error("description ends with a newline");
            }
            constexpr int kMaxLineSize = 78;
            int cur_line_size = 0;
            for (const auto c: s) {
                cur_line_size = (c == '\n') ? 0 : cur_line_size + 1;
                if (cur_line_size > kMaxLineSize) {
                    return turbo::invalid_argument_error("description line length exceeds ", kMaxLineSize,
                                                         " characters");
                }
            }
            return turbo::OkStatus();
        }

    }  // namespace

    turbo::Status Function::validate() const {
        if (!doc_.summary.empty()) {
            // Documentation given, check its contents
            int arg_count = static_cast<int>(doc_.arg_names.size());
            // Some varargs functions allow 0 vararg, others expect at least 1,
            // hence the two possible values below.
            bool arg_count_match = (arg_count == arity_.num_args) ||
                                   (arity_.is_varargs && arg_count == arity_.num_args + 1);
            if (!arg_count_match) {
                return turbo::invalid_argument_error(
                        "In function '", name_,
                        "': ", "number of argument names for function documentation != function arity");
            }
            turbo::Status st = ValidateFunctionSummary(doc_.summary);
            if (st.ok()) {
                st &= ValidateFunctionDescription(doc_.description);
            }
            if (!st.ok()) {
                return st.with_message("In function '", name_, "': ", st.message());
            }
        }
        return turbo::OkStatus();
    }

    turbo::Status ScalarFunction::add_kernel(std::vector<InputType> in_types, OutputType out_type,
                                            ArrayKernelExec exec, KernelInit init) {
        TURBO_RETURN_NOT_OK(check_arity(in_types.size()));

        if (arity_.is_varargs && in_types.size() != 1) {
            return turbo::invalid_argument_error("VarArgs signatures must have exactly one input type");
        }
        auto sig =
                KernelSignature::create(std::move(in_types), std::move(out_type), arity_.is_varargs);
        kernels_.emplace_back(std::move(sig), exec, init);
        return turbo::OkStatus();
    }

    turbo::Status ScalarFunction::add_kernel(ScalarKernel kernel) {
        TURBO_RETURN_NOT_OK(check_arity(kernel.signature->in_types().size()));
        if (arity_.is_varargs && !kernel.signature->is_varargs()) {
            return turbo::invalid_argument_error("Function accepts varargs but kernel signature does not");
        }
        kernels_.emplace_back(std::move(kernel));
        return turbo::OkStatus();
    }

    turbo::Status VectorFunction::add_kernel(std::vector<InputType> in_types, OutputType out_type,
                                            ArrayKernelExec exec, KernelInit init) {
        TURBO_RETURN_NOT_OK(check_arity(in_types.size()));

        if (arity_.is_varargs && in_types.size() != 1) {
            return turbo::invalid_argument_error("VarArgs signatures must have exactly one input type");
        }
        auto sig =
                KernelSignature::create(std::move(in_types), std::move(out_type), arity_.is_varargs);
        kernels_.emplace_back(std::move(sig), exec, init);
        return turbo::OkStatus();
    }

    turbo::Status VectorFunction::add_kernel(VectorKernel kernel) {
        TURBO_RETURN_NOT_OK(check_arity(kernel.signature->in_types().size()));
        if (arity_.is_varargs && !kernel.signature->is_varargs()) {
            return turbo::invalid_argument_error("Function accepts varargs but kernel signature does not");
        }
        kernels_.emplace_back(std::move(kernel));
        return turbo::OkStatus();
    }

    turbo::Status ScalarAggregateFunction::add_kernel(ScalarAggregateKernel kernel) {
        TURBO_RETURN_NOT_OK(check_arity(kernel.signature->in_types().size()));
        if (arity_.is_varargs && !kernel.signature->is_varargs()) {
            return turbo::invalid_argument_error("Function accepts varargs but kernel signature does not");
        }
        kernels_.emplace_back(std::move(kernel));
        return turbo::OkStatus();
    }

    turbo::Status HashAggregateFunction::add_kernel(HashAggregateKernel kernel) {
        TURBO_RETURN_NOT_OK(check_arity(kernel.signature->in_types().size()));
        if (arity_.is_varargs && !kernel.signature->is_varargs()) {
            return turbo::invalid_argument_error("Function accepts varargs but kernel signature does not");
        }
        kernels_.emplace_back(std::move(kernel));
        return turbo::OkStatus();
    }

    turbo::Result<Datum> MetaFunction::execute(const std::vector<Datum> &args,
                                               const FunctionOptions *options,
                                               ExecContext *ctx) const {
        TURBO_RETURN_NOT_OK(CheckArityImpl(*this, static_cast<int>(args.size())));
        TURBO_RETURN_NOT_OK(CheckOptions(*this, options));

        if (options == nullptr) {
            options = default_options();
        }
        return ExecuteImpl(args, options, ctx);
    }

    turbo::Result<Datum> MetaFunction::execute(const ExecBatch &batch,
                                               const FunctionOptions *options,
                                               ExecContext *ctx) const {
        return execute(batch.values, options, ctx);
    }

}  // namespace nebula::compute
